package router

import (
	"context"
	"errors"
	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/schema"
)

const Empty = "DEFAULT"

var ErrNotHandles = errors.New("不存在合适的handler")

type Router struct {
	handlers     map[string]Handler
	chain        chains.Chain
	callbacks    callbacks.Handler
	memory       schema.Memory
	outputparser outputparser.Structured
	emptyHandle  Handler
}

func NewRouter(llm llms.Model, handler []Handler, opts ...Option) *Router {
	opt := executorDefaultOptions(handler)
	for _, o := range opts {
		o(&opt)
	}

	hs := make(map[string]Handler, len(handler))
	for _, h := range handler {
		hs[h.Name()] = h
	}

	return &Router{
		handlers:     hs,
		chain:        chains.NewLLMChain(llm, opt.prompt),
		callbacks:    opt.callback,
		memory:       opt.memory,
		emptyHandle:  opt.emptyHandler,
		outputparser: _outputparser,
	}
}

func (r *Router) Call(ctx context.Context, inputs map[string]any, options ...chains.ChainCallOption) (map[string]any, error) {
	if r.callbacks != nil {
		r.callbacks.HandleChainStart(ctx, inputs)
	}

	if len(r.handlers) == 0 {
		if r.emptyHandle != nil {
			return chains.Call(ctx, r.emptyHandle.Chains(), inputs)
		} else {
			return nil, ErrNotHandles
		}
	}

	result, err := chains.Call(ctx, r.chain, inputs, options...)
	if err != nil {
		return nil, err
	}

	text, ok := result["text"]
	if !ok {
		return nil, chains.ErrNotFound
	}

	out, err := r.outputparser.Parse(text.(string))
	if err != nil {
		return nil, err
	}

	if r.callbacks != nil {
		r.callbacks.HandleChainEnd(ctx, map[string]any{
			"out": out,
		})
	}

	data := out.(map[string]string)
	next, ok := data[_destinations]
	if !ok || next == Empty || r.handlers[next] == nil {
		if r.emptyHandle != nil {
			return chains.Call(ctx, r.emptyHandle.Chains(), inputs)
		} else {
			return nil, ErrNotHandles
		}
	}

	return chains.Call(ctx, r.handlers[next].Chains(), inputs)
}

func (r *Router) GetMemory() schema.Memory {
	return r.memory
}

func (r *Router) GetInputKeys() []string {
	return nil
}

func (r *Router) GetOutputKeys() []string {
	return nil
}
